package lcof2

import (
	"container/heap"

	"gitee.com/ombak/leetcode/pkg/list"
)

type NodeQueue []*list.ListNode

func (n NodeQueue) Less(i, j int) bool  { return n[i].Val < n[j].Val }
func (n NodeQueue) Swap(i, j int)       { n[i], n[j] = n[j], n[i] }
func (n NodeQueue) Len() int            { return len(n) }
func (n *NodeQueue) Push(x interface{}) { *n = append(*n, x.(*list.ListNode)) }
func (n *NodeQueue) Pop() interface{} {
	old := *n
	length := len(old)
	x := old[length-1]
	*n = old[0 : length-1]
	return x
}

// https://leetcode.cn/problems/vvXgSW/
func mergeKLists(lists []*list.ListNode) *list.ListNode {
	dummy := &list.ListNode{}
	p := dummy
	pq := &NodeQueue{}
	heap.Init(pq)
	for _, e := range lists {
		if e != nil {
			heap.Push(pq, e)
		}
	}
	for pq.Len() > 0 {
		x := heap.Pop(pq).(*list.ListNode)
		if x.Next != nil {
			heap.Push(pq, x.Next)
		}
		p.Next = x
		p = p.Next
	}
	p.Next = nil
	return dummy.Next
}
